import os
import os.path as osp


def makedir(path):
    if not osp.exists(path):
        print(f'the {path} is not exists, we mkdir the path.')
        os.makedirs(path)


if __name__=="__main__":
    base_path = 'C:/deep/03_vessel_reid/imgs'
    meta_path = 'C:/deep/03_vessel_reid/meta'
    val_split = 0.2
    makedir(meta_path)
    image_dirs = os.listdir(base_path)
    image_dirs.sort()
    num_ids = len(image_dirs)
    num_train_ids = int(num_ids * (1 - val_split))

    reid_train_list = []
    reid_val_list = []
    train_label, val_label = 0, 0

    reid_entire_dataset_list = reid_train_list.copy()
    for image_dir in image_dirs[:num_train_ids]:
        images = os.listdir(osp.join(base_path, image_dir))
        images.sort()
        for image in images:
            reid_train_list.append(
                f'{image_dir}/{image} {train_label}\n')
            reid_entire_dataset_list.append(f'{image_dir}/{image} {train_label+val_label}\n')
        train_label += 1
    for image_dir in image_dirs[num_train_ids:]:
        images = os.listdir(osp.join(base_path, image_dir))
        images.sort()
        for image in images:
            reid_val_list.append(
                f'{image_dir}/{image} {train_label}\n')
            reid_entire_dataset_list.append(f'{image_dir}/{image} {train_label+val_label}\n')
        val_label += 1
    print(reid_train_list)
    with open(
            osp.join(meta_path,
                     f'train_{int(100 * (1 - val_split))}.txt'),
            'w') as f:
        f.writelines(reid_train_list)
    with open(
            osp.join(meta_path, f'val_{int(100 * val_split)}.txt'),
            'w') as f:
        f.writelines(reid_val_list)
    with open(osp.join(meta_path, 'test.txt'), 'w') as f:
        f.writelines(reid_entire_dataset_list)
